#include "my_pool.h"

static int
my_pool_block_node_cmp(my_rbtree_node_t *a, my_rbtree_node_t *b)
{
    my_pool_block_node_t *a_node, *b_node;
    a_node = (my_pool_block_node_t *)a;
    b_node = (my_pool_block_node_t *)b;

    if (a_node->ptr == b_node->ptr) {
         return 0;
    }

    return a_node->ptr > b_node->ptr ? 1 : -1;
}

my_pool_page_t *
my_pool_page_alloc()
{
    void *ptr = my_thread_malloc(MY_POOL_PAGE_SIZE);
    bzero(ptr, sizeof(my_pool_page_t));
    my_pool_page_t *page = (my_pool_page_t *)ptr;
    page->start = ptr + sizeof(my_pool_page_t);
    page->pos = page->start;
    page->end = ptr + MY_POOL_PAGE_SIZE;
    return page;
}

void
my_pool_page_free(my_pool_page_t *page)
{
    my_thread_free(page);
}

my_pool_cleanup_node_t *
my_pool_add_cleanup(my_pool_t *pool, my_pool_cleanup_pt cleanup, void *data)
{
    my_pool_cleanup_node_t *cleanup_node = my_thread_calloc(sizeof(my_pool_cleanup_node_t));
    cleanup_node->cleanup = cleanup;
    cleanup_node->data = data;
    my_queue_insert_head(&pool->cleanup_queue, &cleanup_node->queue);
    return cleanup_node;
}

void
my_pool_del_cleanup(my_pool_t *pool, my_pool_cleanup_node_t *cleanup_node)
{
    (void)pool;
    my_queue_remove(&cleanup_node->queue);
    my_thread_free(cleanup_node);
}

void *
my_pool_block_alloc(my_pool_t *pool, int64_t size)
{
    void *ptr = my_thread_malloc(size);
    my_pool_block_node_t *block_node = my_thread_calloc(sizeof(my_pool_block_node_t));
    block_node->ptr = ptr;
    my_rbtree_insert(&pool->block_tree, &block_node->tree_node);
    return ptr;
}

void
my_pool_block_free(my_pool_t *pool, void *ptr)
{
    my_pool_block_node_t target = {.ptr = ptr};
    my_rbtree_node_t *node = my_rbtree_lookup(&pool->block_tree, &target.tree_node);
    if (!node) {
        return;
    }

    my_rbtree_delete(&pool->block_tree, node);
    my_pool_block_node_t *block_node = (my_pool_block_node_t *)node;
    my_thread_free(block_node->ptr);
    my_thread_free(block_node);
}

void *
my_pool_mem_alloc(my_pool_t *pool, int64_t size)
{
    size = MY_ALIGN(size, MY_POOL_ALLOC_ALIGNED_SIZE);
    if (size > MY_PAGE_MAX_ALLOC_SIZE) {
        return my_pool_block_alloc(pool, size);
    }

    my_pool_page_t *page = pool->cur_page;
    if (page == NULL || page->end - page->pos < size) {
        page = my_pool_page_alloc();
        my_queue_insert_head(&pool->pages, &page->queue);
    }

    void *ptr = page->pos;
    page->pos += size;

    return ptr;
}

void
my_pool_mem_free(my_pool_t *pool, void *ptr)
{
    my_pool_block_free(pool, ptr);
}

my_pool_t *
my_pool_create()
{
    my_pool_t *pool = my_thread_calloc(sizeof(my_pool_t));
    my_queue_init(&pool->pages);
    my_rbtree_init(&pool->block_tree, &pool->block_sentinel, my_pool_block_node_cmp);
    my_queue_init(&pool->cleanup_queue);

    return pool;
}

void
my_pool_destroy(my_pool_t *pool)
{
    my_queue_t              *q;
    my_pool_cleanup_node_t *cleanup_node;
    while (!my_queue_empty(&pool->cleanup_queue)) {
        q = my_queue_last(&pool->cleanup_queue);
        cleanup_node = my_queue_data(q, my_pool_cleanup_node_t, queue);
        cleanup_node->cleanup(pool, cleanup_node->data);
        my_pool_del_cleanup(pool, cleanup_node);
    }

    my_rbtree_node_t *node;
    my_pool_block_node_t *block_node;
    while(!my_rbtree_empty(&pool->block_tree)) {
        node = my_rbtree_min(pool->block_tree.root, pool->block_tree.sentinel);
        block_node = (my_pool_block_node_t *)node;
        my_pool_block_free(pool, block_node->ptr);
    }

    my_pool_page_t *page;
    while (!my_queue_empty(&pool->pages)) {
        q = my_queue_last(&pool->pages);
        my_queue_remove(q);
        page = my_queue_data(q, my_pool_page_t, queue);
        my_pool_page_free(page);
    }

    my_thread_free(pool);
}